from functools import cached_property
import re
import argparse
import logging
from pathlib import Path
import polars as pl
from pydantic import BaseModel, computed_field

logger = logging.getLogger("BioMedDB Command")


class PrepareRepTrans(BaseModel):
    ncbi_refseq_curated_file: Path
    ncbi_refseq_hgmd_file: Path
    ncbi_mane_select_file: Path
    rep_trans_file: Path
    ncbi_refseq_representative_file: Path

    @classmethod
    def format_mane_chrom(cls, chromosome: str):
        chromosome = chromosome.split(".")[0][7:9].lstrip("0")
        return "chr" + {"23": "X", "24": "Y"}.get(chromosome, chromosome)

    @cached_property
    def mane_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_mane_select_file)
        records = pl.read_csv(
            self.ncbi_mane_select_file,
            sep="\t",
            columns=["symbol", "RefSeq_nuc", "MANE_status", "GRCh38_chr"],
            new_columns={"gene": str, "transcript": str, "status": str, "chromosome": str},
        )
        records = records.filter(pl.col("chromosome").str.startswith("NC_")).with_columns(
            pl.col("chromosome").map_elements(self.format_mane_chrom, pl.String),
            pl.col("transcript").str.split(".").list.first().alias("trans"),
            pl.when(pl.col("status").str.to_lowercase().str.contains("clinical")).then(pl.lit("clinical")).otherwise(pl.col("mane")).alias("source"),
        )
        return records.select(["chromosome", "gene", "trans", "source"])

    @cached_property
    def hgmd_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_refseq_hgmd_file)
        records = pl.read_csv(
            self.ncbi_refseq_hgmd_file,
            separator="\t",
            has_header=False,
            columns=[1, 2, 12],
            new_columns={"transcript": str, "chromosome": str, "gene": str},
        )
        records = records.filter(pl.col("chromosome").str.len_chars().le(5)).with_columns(
            pl.col("transcript").str.split(".").list.first().alias("trans"),
            pl.lit("hgmd").alias("source"),
        )
        return records.select(["chromosome", "gene", "trans", "source"])

    @cached_property
    def max_size_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_refseq_curated_file)
        records = pl.read_csv(
            self.ncbi_refseq_curated_file,
            separator="\t",
            has_header=False,
            columns=[1, 2, 4, 5, 6, 7, 12],
            new_columns={"transcript": str, "chromosome": str, "tx_start": int, "tx_end": int, "cds_start": int, "cds_end": int, "gene": str},
        )
        records = records.filter(pl.col("chromosome").str.len_chars().le(5)).with_columns(
            pl.col("transcript").str.split(".").list.first().alias("trans"),
            (pl.col("tx_end") - pl.col("tx_start")).alias("tx_size"),
            (pl.col("cds_end") - pl.col("cds_start")).alias("cds_size"),
            pl.lit("max_size").alias("source"),
        )
        records = records.group_by(["chromosome", "gene"]).agg(pl.all().filter(pl.col("tx_size").eq(pl.col("tx_size").max())).head(1)).explode(records.columns)
        return records.select(["chromosome", "gene", "trans", "source"])

    @cached_property
    def rep_trans_records(self) -> pl.DataFrame:
        keys = ["chromosome", "gene", "trans"]
        records = pl.read_csv(self.ncbi_refseq_curated_file, separator="\t", has_header=False)
        records = records.rename({"column_2": "transcript", "column_3": "chromosome", "column_13": "gene"})
        records = records.with_columns(pl.col("transcript").str.split(".").list.first().alias("trans"))
        # 首次注释来源
        records = records.join(self.mane_records, on=keys, how="left")
        # 合并HGMD来源
        records = records.join(self.hgmd_records, on=keys, how="left").with_columns(pl.col("source").fill_null(pl.col("source_right"))).drop("source_right")
        # 合并最大长度来源
        records = records.join(self.max_size_records, on=keys, how="left").with_columns(pl.col("source").fill_null(pl.col("source_right"))).drop("source_right")
        return records.drop("trans")

    def run(self):
        records = self.rep_trans_records
        records.select(["chromosome", "gene", "transcript", "source"]).write_csv(self.rep_trans_file, separator="\t")
        records.drop("source").write_csv(self.ncbi_refseq_representative_file, separator="\t", include_header=False)


class PrepareGeneID(BaseModel):
    ncbi_refseq_curated_file: Path
    ncbi_gene2refseq_file: Path
    ncbi_gene_info_file: Path
    gene_id_file: Path

    @cached_property
    def gene2refseq_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_gene2refseq_file)
        records = pl.read_csv(
            self.ncbi_gene2refseq_file,
            separator="\t",
            columns=["#tax_id", "GeneID", "RNA_nucleotide_accession.version", "genomic_nucleotide_accession.version"],
            new_columns={"tax_id": str, "gene_id": str, "transcript": str, "chromosome": str},
        )
        records = records.filter(pl.col("tax_id").eq("9606")).with_columns(pl.col("transcript").str.split(".").list.first().alias("trans"))
        records = records.with_columns(pl.col("chromosome").map_elements(PrepareRepTrans.format_mane_chrom, pl.String))
        return records.select(["chromosome", "trans", "gene_id"]).unique(["chromosome", "trans"])

    @cached_property
    def gene_info_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_gene_info_file)
        records = pl.read_csv(
            self.ncbi_gene_info_file,
            separator="\t",
            columns=["GeneID", "Symbol", "Synonyms", "chromosome"],
            new_columns={"gene_id": str, "symbol": str, "synonyms": str, "chromosome": str},
            null_values=["-", "."],
        )
        logger.info("Processing NCBI Gene Info")
        records = (
            records.with_columns(pl.col("chromosome").map_elements(lambda x: [f"chr{c}" for c in re.split(r"\W+", x)], pl.List(pl.String)).alias("chromosome"))
            .explode("chromosome")
            .with_columns(pl.col("synonyms").str.split("|").alias("synonyms"))
        )
        # symbol_records = records.select(["chromosome", "symbol", "gene_id"]).rename({"symbol": "gene"}).unique(["chromosome", "gene"])
        # synonyms_records = records.select(["chromosome", "synonyms", "gene_id"]).rename({"synonyms": "gene"}).unique(["chromosome", "gene"])
        return records

    @cached_property
    def genepred_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_refseq_curated_file)
        records = pl.read_csv(
            self.ncbi_refseq_curated_file,
            separator="\t",
            has_header=False,
            usecols=[1, 2, 12],
            columns=["transcript", "chromosome", "gene"],
        )
        records = (
            records.filter(pl.col("chromosome").str.len_chars().le(5))
            .with_columns(pl.col("transcript").str.split(".").list.first().alias("trans"))
            .unique(["chromosome", "gene"])
            .sort(["chromosome", "gene"])
        )
        return records.select(["chromosome", "gene", "trans"])

    def run(self) -> pl.DataFrame:
        symbol_records = self.gene_info_records.select(["chromosome", "symbol", "gene_id"]).rename({"symbol": "gene"}).unique(["chromosome", "gene"])
        synonyms_records = self.gene_info_records.select(["chromosome", "synonyms", "gene_id"]).explode("synonyms").rename({"symbol": "gene"}).unique(["chromosome", "gene"])
        # 初始化GeneID
        records = self.genepred_records.join(self.gene2refseq_records, on=["chromosome", "trans"])
        # 合并Symbol GeneID
        records = records.join(symbol_records, on=["chromosome", "gene"], how="left").with_columns(pl.col("gene_id").fill_null(pl.col("gene_id_right"))).drop("gene_id_right")
        # 合并Synonyms GeneID
        records = records.join(synonyms_records, on=["chromosome", "gene"], how="left").with_columns(pl.col("gene_id").fill_null(pl.col("gene_id_right"))).drop("gene_id_right")

        logger.info("Writing to %s", self.gene_id_file)
        records.select(["chromosome", "gene", "gene_id"]).write_csv(self.gene_id_file, separator="\t")


class PrepareRefseq(BaseModel):
    ncbi_refseq_curated_file: Path
    ncbi_refseq_hgmd_file: Path
    ncbi_mane_select_file: Path
    ncbi_gene2refseq_file: Path
    ncbi_gene_info_file: Path
    outdir: Path

    @property
    @computed_field
    def rep_trans_file(self) -> Path:
        return self.outdir / "representative_transcript.txt"

    @property
    @computed_field
    def gene_id_file(self) -> Path:
        return self.outdir / "gene_id.txt"

    @property
    @computed_field
    def ncbi_refseq_representative_file(self) -> Path:
        return self.outdir / "ncbiRefSeqRepresentative.txt"

    def run(self):
        self.outdir.mkdir(parents=True, exist_ok=True)
        kwargs = self.model_dump()
        PrepareRepTrans.model_validate(kwargs).run()
        PrepareGeneID.model_validate(kwargs).run()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare Refseq")
    parser.add_argument("--ncbi_refseq_curated_file", "-curated", type=Path, required=True, help="input, ncbi refseq curated_file")
    parser.add_argument("--ncbi_refseq_hgmd_file", "-hmgd", type=Path, required=True, help="input, ncbi refseq hgmd file")
    parser.add_argument("--ncbi_mane_select_file", "-mane", type=Path, required=True, help="input, ncbi mane select file")
    parser.add_argument("--ncbi_gene2refseq_file", "-gene2refseq", type=Path, required=True, help="input, ncbi gene2refseq select file")
    parser.add_argument("--ncbi_gene_info_file", "-gene_info", type=Path, required=True, help="input, ncbi gene info select file")
    parser.add_argument("--outdir", "-out", type=Path, required=True, help="outdir, outdir")
    parser.set_defaults(func=lambda args: PrepareRefseq.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
